%% Data
clear all

addpath('...\replication\Section 6\01_data');

for uc = [1 4];
    
c = 12;
    
%% Read original data
L_ni = dlmread('...\replication\Section 6\01_data\commutes_ij_final.csv'); 
d_ni_raw = dlmread('...\replication\Section 6\01_data\trav_dist_ij_final.csv'); 
n_vars = dlmread('...\replication\Section 6\01_data\i_vars.csv'); 
tau_ni = dlmread('...\replication\Section 6\01_data\tau_ij.csv'); 

%% Consistent with commutes (i.e., all locs are orig and dest)
L_i = sum(L_ni);
L_i_p = (L_i~=0);

L_ni = L_ni(L_i_p',L_i_p);
d_ni_raw = d_ni_raw(L_i_p',L_i_p);
n_vars = n_vars(L_i_p',:);

%% Set vars after corrections
R_n = sum(L_ni,2);
L_i = sum(L_ni);

com = n_vars(:,1);
csvwrite('com_final.csv',com);
H_R = n_vars(:,5);
Q_n = n_vars(:,6).*(90.9*12);
v_n = n_vars(:,7);
v_n=v_n./v_n(1:1);
xy_i = [n_vars(:,2),n_vars(:,3)];

%% Parameters

L=sum(sum(L_ni),2);    % Total population 
alpha=0.8 ;          
xi = 1.34;       % semi-elasticity of commutes to distance
phi = 0.43;            % semi-elasticity of trade to distance
sigma = 4;
epsilon = 2;
n=length(R_n);    % size of locations
eta = 0.9;
tau = 4.7;



%% Read cal data
% Prod and amenities
B_n = dlmread('B_n.csv'); 
A_i = dlmread('A_i.csv'); 
% Res inc
w_i_raw = dlmread('w_i.csv');
w_i=w_i_raw./geomean(w_i_raw);


%% Commutes and trade counterfactrual scalars
d_ni_tilde = pdist2(xy_i,xy_i,'euclidean');
d_ni_tilde = 100.*d_ni_tilde; % Make km
d_ni_tilde(d_ni_tilde == 0) = 1;



%% Gen scenarios

if uc == 1
    unit_cost = 0; 
    d_ni_tilde = tau.*tau_ni + d_ni_tilde;
elseif uc == 2
    unit_cost = 2000000 + 400000; 
elseif uc == 3
    unit_cost = 3000000 + 400000; 
elseif uc == 4
    unit_cost = 4000000 + 400000; 
else uc == 4
    unit_cost = 5000000 + 400000;
end

diff_km = 4588.748; 
cost_for_each_km = diff_km.*unit_cost;
cost = cost_for_each_km; 
t = cost/L;

d_ni = tau.*tau_ni + d_ni_tilde;

% Kappa
kappa = (d_ni_tilde./d_ni).^(xi);
% Trade
zeta = (d_ni_tilde./d_ni).^(-phi);

%% Remaining vars

lambda_ni = L_ni./L;

num = L_i.*(zeta.*w_i./A_i).^(1-sigma);
pi_ni = num./sum(num);

Y_i = sum(w_i.*lambda_ni.*L);
Y_n = Y_i';

%% Solve modified Monte et al (2018) (see Appendix - Section B.2)

%Initial values
w_i0= w_i;
lambda_ni0=ones(n,n);

error1=10^-3;
error2=10^-3;

it = 1;
while it < 1000
%Step 1
num1 = B_n.*lambda_ni.*((w_i0-t)./kappa).^epsilon;
v_n0 = (1./v_n).*(sum(num1.*(w_i0-t).*(w_i-t)./sum(num1,2) ,2));

L_i0 = (L./L_i).*sum(lambda_ni.*lambda_ni0);

R_n0 = (L./R_n).*sum(lambda_ni.*lambda_ni0,2);

%Step 2:
Q_n0 = (v_n0.*R_n0).^(1-eta);

%Step 3:
num2 = L_i0.*(zeta.*w_i0./A_i).^(1-sigma);
pi_ni0 = num2./sum(pi_ni.*num2,2);

%Step 4:
pi_nn0 = diag(pi_ni0);
P_n0 = ((L_i0'./pi_nn0).^(1/(1-sigma))).*w_i0'./A_i';

%Step 5:
w_i1 = ((1./(Y_i.*L_i0)).*(sum(pi_ni.*pi_ni0.*v_n0.*R_n0.*Y_n))) + t;
w_i1=w_i1./geomean(w_i1);

num3 = B_n.*(P_n0.^alpha .* Q_n0.^(1-alpha)).^(-epsilon) .* (w_i0./kappa).^epsilon;
lambda_ni1 = num3./sum(sum(num3,2));

%Step 6:
% Stopping criteria
val1 = max(abs(w_i1-w_i0));
val2 = max(max(abs(lambda_ni1-lambda_ni0)),[],2);
if (val1 < error1) && (val2 < error2)
it = 10^8;
display('>>>> System Converged <<<<');
w_i_final=w_i1;
lambda_ni_final = lambda_ni1;
P_n_final = P_n0;
Q_n_final = Q_n0;

%Update criteria
else
    
w_i0 =  0.3.*w_i1 + 0.7.*w_i0;
lambda_ni0 = 0.3.*lambda_ni1 + 0.7.*lambda_ni0;

        % Choose units for w_i such that geometric mean w_i=1;
  t = (cost/L)/geomean(w_i0);
  w_i0=w_i0./geomean(w_i0);
end
it = it + 1;
end

% Real res income
lambda_ni_n = lambda_ni_final./sum(lambda_ni_final,2);
v_n_final = sum(lambda_ni_n.*w_i_final,2);
real_res_inc = v_n_final./((Q_n_final.^(1-alpha)).*(P_n_final.^(alpha)));

% Residential population
R_n_final = sum(lambda_ni_final.*L,2);

% Total commutes
com_n = sum(d_ni.*lambda_ni_final.*L,2)./R_n_final;
tot_com = sum(d_ni.*lambda_ni_final.*L,2);

% Total res income
Y_n = v_n_final.*R_n_final;

% Sum of total real res income
res_Y =  sum(real_res_inc.*sum(lambda_ni_final.*L,2))

% Workplace employment
L_i_final = sum(lambda_ni_final.*L);

% Total workplace income
Y_i = w_i_final.*L_i_final;

res(uc,:) = [c,res_Y];

res_Y_n(:,uc) = [Y_n];
res_Y_i(:,uc) = [Y_i];

res_w_i(:,uc) = [w_i_final];
res_v_n(:,uc) = [v_n_final];

res_L_i(:,uc) = [L_i_final];
res_R_n(:,uc) = [R_n_final];

res_com_n(:,uc) = [com_n];
res_Tcom_n(:,uc) = [tot_com];
res_Q_n(:,uc) = [Q_n_final];

clearvars -except res res_Y_n res_Y_i res_L_i res_R_n uc res_w_i res_v_n res_com_n res_Q_n res_Tcom_n

end

% Export data
dlmwrite('...\replication\Section 6\01_data\res_residential_inc.csv',res,'precision',16);

dlmwrite('...\replication\Section 6\01_data\res_Y_n.csv',res_Y_n,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_Y_i.csv',res_Y_i,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_L_i.csv',res_L_i,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_R_n.csv',res_R_n,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_w_i.csv',res_w_i,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_v_n.csv',res_v_n,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_com_n.csv',res_com_n,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_Tcom_n.csv',res_Tcom_n,'precision',16);
dlmwrite('...\replication\Section 6\01_data\res_Q_n.csv',res_Q_n,'precision',16);

